from typing import Sequence

import numpy as np


class Light:
    GLSL_NAME = 'lights'

    def __init__(self):
        self.color = (1.0, 1.0, 1.0)
        self.strength = 1.0

    def __iter__(self):
        yield self.color[0]
        yield self.color[1]
        yield self.color[2]
        yield self.strength


class StraightLight(Light):
    GLSL_NAME = 'straightLights'

    def __init__(self):
        super().__init__()
        self.direction = (0, -1, 0)

    def __iter__(self):
        yield self.color[0]
        yield self.color[1]
        yield self.color[2]
        yield self.direction[0]
        yield self.direction[1]
        yield self.direction[2]
        yield self.strength


class PointLight(Light):
    GLSL_NAME = 'pointLights'

    def __init__(self):
        super().__init__()
        self.position = (0, 0, 0)

    def __iter__(self):
        yield self.color[0]
        yield self.color[1]
        yield self.color[2]
        yield self.position[0]
        yield self.position[1]
        yield self.position[2]
        yield self.strength


class LightManager:
    def __init__(self):
        self.lights = set()

    def add_point_light(self, color: Sequence[float], position: Sequence[float], strength: float = 1.0):
        point_light = PointLight()
        point_light.color = color
        point_light.position = position
        point_light.strength = strength
        self.lights.add(point_light)
        return point_light

    def add_straight_light(self, color: Sequence[float], direction: Sequence[float], strength: float = 1.0):
        straight_light = StraightLight()
        straight_light.color = color
        straight_light.direction = direction
        straight_light.strength = strength
        self.lights.add(straight_light)
        return straight_light

    def add_light(self, color: Sequence[float], strength: float = 1.0):
        light = Light()
        light.color = color
        light.strength = strength
        self.lights.add(light)
        return light

    def remove_light(self, light: Light):
        if light in self.lights:
            self.lights.remove(light)

    @property
    def uniforms(self):
        data: dict = {}
        size: dict = {}
        for light in self.lights:
            glsl_name = light.GLSL_NAME
            if glsl_name not in data:
                data[glsl_name] = []
                size[glsl_name] = 0
            data[glsl_name].extend([*light])
            size[glsl_name] += 1
        for key, value in data.items():
            yield f'{key}.data', np.array(value, dtype='f4')
        for key, value in size.items():
            yield f'{key}.count', value
        for key in size.keys():
            yield f'has{key}', True
